'''
Created on Feb 26, 2010

@author: mkiyer
'''
# built-in python imports
import collections
import os
import pickle
import logging
import subprocess
import shutil

# globally installed packages
import ruffus

# project imports
import veggie.alignment.sam as sam
from veggie.genome.genome import get_genome

# local imports
import pipelineconfig
from chipseqdb import ChIPSeqDB

class SamplePoolTask(object):
    def __init__(self):
        self.db_connection_string = None
        self.db_schema_name = None
        self.sample_pool_id = None
        self.genome_version = None
        self.fragment_length = None
        self.make_norm_bigwig = False
        self.make_raw_bigwig = False
        self.output_dir = None
        self.output_file_map = None
        self.library_ids = []
        self.alignment_files = []
        self.ucsc_bin_dir = None
        self.chrom_sizes_file = None

def update_status_complete(task_info):
    logger = logging.getLogger(task_info.sample_pool_id)
    status_flag = ChIPSeqDB.SAMPLE_POOL_STATUS_COMPLETE
    logger.debug('Marking sample pool status_flag=%s' % (status_flag))
    db = ChIPSeqDB(task_info.db_connection_string, 
                   task_info.db_schema_name, echo=False)    
    db.set_sample_pool_status_flag(task_info.sample_pool_id, status_flag)

def filter_references(insamfile, outsamfile, references):    
    if (references == None):
        return
    logger = logging.getLogger(__name__)
    outsam = open(outsamfile, "w")
    insam = open(insamfile, "r")
    for header_line in insam:
        # once we find the first line that does not start with '@' the header is over
        if not header_line.startswith("@"):
            break
        header_fields = header_line.strip().split('\t')
        if header_fields[0] == '@SQ':
            fields = dict([field.split(':', 1) for field in header_fields[1:]])
            assert 'SN' in fields
            if fields['SN'] not in references:
                logger.debug('filter_references: rejected reference %s' % fields['SN'])
                continue        
        outsam.write(header_line)
    # still have to deal with the first read before we can start a new loop
    rejected_reads = 0    
    getrname = lambda r: r.split('\t', 3)[2]
    read_line = header_line
    rname = getrname(read_line)
    if rname == '*':
        outsam.write(read_line)
    elif rname in references:
        outsam.write(read_line)
    else:
        rejected_reads += 1
    # loop through rest of reads
    for read_line in insam:
        rname = getrname(read_line)
        if rname == '*':
            outsam.write(read_line)
        elif rname in references:
            outsam.write(read_line)
        else:
            rejected_reads += 1    
    insam.close()
    outsam.close()
    logger.debug('filter_references: removed %d reads from non-matching references' % (rejected_reads))
    return outsamfile

def convert_sam_to_bam(insamfiles, bamfile, references, tmp_path=None):
    logger = logging.getLogger(__name__)
    tmpfiles = []
    mergebamfiles = []
    if tmp_path == None:
        tmp_path = os.path.abspath(os.path.dirname(bamfile))
    # first have to convert the SAM files to BAM and prepare them
    for insamfile in insamfiles:
        # remove references that are contamination
        myfiltersam = os.path.join(tmp_path, os.path.splitext(os.path.basename(insamfile))[0] + '.filter.sam')
        tmpfiles.append(myfiltersam)
        logger.debug("Filtering references from SAM %s -> SAM %s" % (insamfile, myfiltersam))
        filter_references(insamfile, myfiltersam, references=references)
        nextinfile = myfiltersam
        # convert to BAM
        mybamfile = os.path.splitext(nextinfile)[0] + '.bam'
        tmpfiles.append(mybamfile)
        logger.debug("Converting SAM %s -> BAM %s" % (nextinfile, mybamfile))
        subprocess.call(['samtools', 'view', '-S', nextinfile, '-b', '-o', mybamfile])
        nextinfile = mybamfile 
        # sort
        mysortedbam = os.path.splitext(nextinfile)[0] + '.sorted.bam'
        tmpfiles.append(mysortedbam)
        logger.debug("Sorting BAM %s -> %s" % (nextinfile, mysortedbam))
        # sort automatically adds ".bam" to the file, so remove the extension
        subprocess.call(['samtools', 'sort', nextinfile, os.path.splitext(mysortedbam)[0]])
        # index        
        logger.debug("Indexing Sorting BAM %s" % (mysortedbam))
        subprocess.call(['samtools', 'index', mysortedbam])
        tmpfiles.append(mysortedbam + '.bai')
        nextinfile = mysortedbam
        # remove duplicates
        myrmdupbam = os.path.splitext(nextinfile)[0] + '.rmdup.bam'
        tmpfiles.append(myrmdupbam)
        logger.debug("Removing duplicates from BAM %s -> %s" % (nextinfile, myrmdupbam))
        subprocess.call(['samtools', 'rmdup', '-s', nextinfile, myrmdupbam])
        nextinfile = myrmdupbam
        # files to merge
        mergebamfiles.append(nextinfile)        
    # merge
    if len(mergebamfiles) == 1:
        # just one file, so copy it rather than merge it
        # it is already sorted too.. so just re-index it
        logger.debug("Copying BAM file %s -> %s" % (mergebamfiles[0], bamfile))
        shutil.copyfile(mergebamfiles[0], bamfile)
    else:
        #mymergedbam = os.path.splitext(bamfile)[0] + '.merge.bam'
        #tmpfiles.append(mymergedbam)
        merge_cmd = ['samtools', 'merge', bamfile]
        merge_cmd.extend(mergebamfiles)
        logger.debug("Merging BAM files %s -> %s" % (mergebamfiles, bamfile))
        subprocess.call(merge_cmd)
        # sort
        #logger.debug("Sorting BAM file %s -> %s" % (mymergedbam, bamfile))
        # sort automatically adds ".bam" to the file, so remove the extension
        #subprocess.call(['samtools', 'sort', mymergedbam, os.path.splitext(bamfile)[0]])
    logger.debug("Indexing BAM file %s" % (bamfile))
    subprocess.call(['samtools', 'index', bamfile])
    # remove temporary files
    for tmpfile in tmpfiles:
        logger.debug("Removing %s" % tmpfile)
        os.remove(tmpfile)

def wigToBigWig(infile, outfile, task_info):
    executable = os.path.join(task_info.ucsc_bin_dir, 'wigToBigWig')
    chrom_sizes_file = task_info.chrom_sizes_file
    cmd = ' '.join([executable, '-clip', infile, chrom_sizes_file, outfile])
    retcode = os.system(cmd)
    return retcode

@ruffus.split(["config.xml"], "*/*/sample_pool_info.pickle")
def generate_sample_pool_tasks(input_files, output_files):
    logger = logging.getLogger(__name__)
    # read configuration
    logger.info("Parsing configuration file...")
    config = pipelineconfig.PipelineConfig(input_files[0])
    # change directories to the base output dir
    # (some paths are relative to current working dir)
    assert os.path.exists(config.output_dir)
    os.chdir(config.output_dir)
    # setup output path
    sample_pools_dir = os.path.join(config.output_dir, config.sample_pools_dir)
    # log sample pool configuration
    logger.info("Alignment results dir: %s" % config.alignment_dir)
    logger.info("Sample pool output dir: %s" % sample_pools_dir)
    logger.info("Force creation of completed sample pools: %s" % config.sample_pools_force)
    logger.info("Normalized bigwigs: %s" % config.make_norm_bigwig)
    logger.info("Raw read bigwigs: %s" % config.make_raw_bigwig)
    logger.info("cDNA fragment length (for tag extension): %s" % config.cdna_fragment_length)

    # fetch the chromosome sizes using UCSC's script
    logger.info("Fetching chrom sizes from UCSC using MYSQL query...")
    chrom_sizes_file = 'ucsc.chrom.sizes'
    executable = os.path.join(config.ucsc_bin_dir, 'fetchChromSizes')
    cmd = ' '.join(['sh', executable, config.genome_version, '>', chrom_sizes_file])
    os.system(cmd) 

    # connect to database
    connection_string, schema_name = config.databases[config.db]    
    logger.info("Connecting to database [%s] schema=[%s] at %s..." % (config.db, schema_name, connection_string))
    db = ChIPSeqDB(connection_string, schema_name, echo=False)

    # determine which sample pools to process
    sample_pools = collections.defaultdict(lambda: set())
    for sp in db.get_sample_pool_main():
        # skip samples that are marked complete
        if ((sp.status_flag == db.SAMPLE_POOL_STATUS_COMPLETE) and
            config.sample_pools_force == False):
            continue
        # get all libraries for this sample pool    
        libraries = db.get_sample_pool_libraries(sp.sample_pool_id) 
        for flowcell_id, lane_id in libraries:
            # TODO: this is a hack to convert a flowcell/lane into a library id
            # for publicly available data this won't work.  We need to use the
            # sample ID to lookup where the alignments are stored
            library_id = '%s_%s' % (flowcell_id, lane_id)        
            alignment_file = os.path.join(config.alignment_dir, library_id + '.sam')
            if os.path.exists(alignment_file):
                sample_pools[sp.sample_pool_id].add((library_id, alignment_file))
            else:
                logger.warning('SP_ID %s: could not find alignment results for library %s at path: %s' % 
                               (sp.sample_pool_id, library_id, alignment_file))
    sample_pools = dict(sample_pools)

    # create tasks for each sample pool
    for sample_pool_id, sample_pool_libraries in sample_pools.iteritems():
        task_info = SamplePoolTask()
        task_info.db_connection_string = connection_string
        task_info.db_schema_name = schema_name
        task_info.sample_pool_id = sample_pool_id
        task_info.genome_version = config.genome_version
        task_info.fragment_length = config.cdna_fragment_length
        task_info.make_norm_bigwig = config.make_norm_bigwig
        task_info.make_raw_bigwig = config.make_raw_bigwig
        task_info.output_dir = os.path.join(sample_pools_dir, sample_pool_id)
        task_info.output_file_map = db.get_sample_pool_output(sample_pool_id)
        for library_info in sample_pool_libraries:
            library_id, alignment_file = library_info
            task_info.library_ids.append(library_id)
            task_info.alignment_files.append(alignment_file)
        # setup a new working directory for each sample pool
        if not os.path.exists(task_info.output_dir):
            logger.debug('%s: creating sample pool results path %s' % (sample_pool_id, task_info.output_dir))
            os.makedirs(task_info.output_dir)
        # copy the chrom sizes file to the task directory
        task_info.ucsc_bin_dir = config.ucsc_bin_dir
        task_info.chrom_sizes_file = 'ucsc.chrom.sizes'
        shutil.copyfile(chrom_sizes_file, os.path.join(task_info.output_dir, task_info.chrom_sizes_file))
        # TODO: instead of using pickle, use something human readable like XML
        task_info_file = os.path.join(task_info.output_dir, 'sample_pool_info.pickle')
        pickle.dump(task_info, open(task_info_file, 'w'))

@ruffus.transform(generate_sample_pool_tasks, 
                  ruffus.regex(r"(.*)/(.*)/sample_pool_info.pickle$"), 
                  r"\1/\2/merged_alignments.bam", r"\2")
def merge_alignments_task(infile, outfile, sample_pool_id):
    logger = logging.getLogger(sample_pool_id)
    task_info = pickle.load(open(infile))
    if len(task_info.alignment_files) == 0:
        return
    if not os.path.exists(outfile):
        mygenome = get_genome(task_info.genome_version)
        chrom_names = mygenome.get_chrom_names()
        convert_sam_to_bam(task_info.alignment_files, outfile, chrom_names)
    else:
        logger.debug('BAM file %s exists, skipping..' % outfile)

@ruffus.transform(merge_alignments_task, 
                  ruffus.regex(r"(.*)/(.*)/merged_alignments.bam$"), 
                  r"\1/\2/sample_pools.done", r"\2")
def bam_to_wiggle_task(input_file, output_file, sample_pool_id):
    logger = logging.getLogger(sample_pool_id)
    if not os.path.exists(input_file):
        logger.error("Input file %s does not exist" % (input_file))
        return
    # get task info
    task_info_file = os.path.join(os.path.dirname(input_file), "sample_pool_info.pickle")
    task_info = pickle.load(open(task_info_file))
    # get output dir
    output_dir = os.path.dirname(output_file)
    # make bigwig files
    params = []
    for sp_output_type, sp_output_file in task_info.output_file_map.iteritems():
        if sp_output_type == ChIPSeqDB.OUTPUT_TYPE_NORMALIZED:
            params.append(('norm.wig.gz', True, sp_output_file))
        elif sp_output_type == ChIPSeqDB.OUTPUT_TYPE_RAW:
            params.append(('raw.wig.gz', False, sp_output_file))
    
    noerror = True
    for param in params:
        wigfilename, norm_flag, bwfilename = param
        wigfile = os.path.join(output_dir, wigfilename)
        if not os.path.exists(wigfile):
            logger.debug("Converting BAM %s -> wiggle %s" % (input_file, wigfile))
            sam.bam_to_wiggle(input_file, wigfile, 
                              unique_only=False, 
                              merge_strands=True, 
                              fragment_length=task_info.fragment_length, 
                              norm=norm_flag)
        # convert to bigwig
        bwfile = os.path.join(output_dir, bwfilename)
        if not os.path.exists(bwfile):
            logger.debug("Converting WIG %s -> BIGWIG %s" % (wigfile, bwfile))
            retcode = wigToBigWig(wigfile, bwfile, task_info)
            if retcode != 0:
                noerror = False
                logger.error("wigToBigWig returned error code %d" % retcode)        

    if noerror == True:
        # write a blank file signifying task finished successfully
        open(output_file, 'w').close()
        # update the database
        update_status_complete(task_info)

if __name__ == '__main__':
    from optparse import OptionParser    
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    # parse command line arguments
    optionparser = OptionParser("usage: %prog [options]")
    optionparser.add_option("-p", "--processors", dest="processors", type="int", default=4,
                            help="process only the samples defined in the sample XML file")                       
    (options, args) = optionparser.parse_args()
    # run the pipeline
    ruffus.pipeline_run([bam_to_wiggle_task], verbose=2, multiprocess=options.processors)    
